{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3tN0poN3zZ4L"
      },
      "source": [
        "# Meta-Learning\n",
        "\n",
        "[![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.sandbox.google.com/github/google-deepmind/optax/blob/main/examples/meta_learning.ipynb)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wGHOAQQNWSvl"
      },
      "source": [
        "Here we investigate using optax to meta-learn the learning rate of an optax optimizer. For a concrete example, we define a model where $y$ is linearly related to $x$, with some added noise,\n",
        "\n",
        "$$y = f(x) = 10 \\cdot x + \\mathcal{N}(0, 1).$$\n",
        "\n",
        "We imagine trying to solve the problem where we have access to some data generated by $f(\\cdot)$, but we do not know what $f$ actually is. Thus, we can attempt to  approximate $f(\\cdot)$ using a parametrized function, $f(\\theta, x).$ We can find a good value of $\\theta$ using gradient descent with optax. \n",
        "\n",
        "Gradient descent typically requires hyper-parameters be chosen, for example the learning rate of the optimizer. Using meta-learning, we are able to simultaneously optimize over parameters of our optimizer, whilst also optimizing over the underlying parameters for the function we are fitting.\n",
        "\n",
        "Many meta-learning algorithms now take the following form,\n",
        "\n",
        "1. An inner loop performs $N$ updates to a model's parameters $\\theta$ using the current value of a hyper-parameter $\\eta$.\n",
        "1. An outer loop updates the hyper-parameter $\\eta$ to maximize the quality of the inner updates to $\\theta$ by differentiating through the inner updates.\n",
        "\n",
        "Since in our case we are interested in the learning rate, which is\n",
        "bounded between 0 and 1, we parametrize the learning rate as a sigmoid\n",
        "over the meta parameter $\\eta$.\n",
        "\n",
        "In the following snippets, we will solve the problem using optax. To begin with, we define a generator that samples from the hidden underlying distribution."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Z6PrJuBOebo1"
      },
      "outputs": [],
      "source": [
        "from collections.abc import Callable\n",
        "from typing import Iterator, Tuple\n",
        "import chex\n",
        "import jax\n",
        "import jax.numpy as jnp\n",
        "import matplotlib.pyplot as plt\n",
        "import numpy as np\n",
        "import optax"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NA3VjroNZjlo"
      },
      "outputs": [],
      "source": [
        "def generator() -> Iterator[Tuple[chex.Array, chex.Array]]:\n",
        "  rng = jax.random.PRNGKey(0)\n",
        "\n",
        "  while True:\n",
        "    rng, k1, k2 = jax.random.split(rng, num=3)\n",
        "    x = jax.random.uniform(k1, minval=0.0, maxval=10.0)\n",
        "    y = 10.0 * x + jax.random.normal(k2)\n",
        "    yield x, y"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lUfPZ8FVfqc9"
      },
      "source": [
        "We demonstrate sampling from this as follows,"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tDpR8GFfftPD"
      },
      "outputs": [],
      "source": [
        "g = generator()\n",
        "\n",
        "for _ in range(5):\n",
        "  x, y = next(g)\n",
        "  print(f\"Sampled y = {y:.3f}, x = {x:.3f}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "J-L8Vyc9eFTy"
      },
      "source": [
        "We now define our parametrized function $f(\\theta, x)$, and choose a random initial value for the parameter $\\theta$."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "N8PNwdASeNZc"
      },
      "outputs": [],
      "source": [
        "def f(theta: chex.Array, x: chex.Array) -> chex.Array:\n",
        "  return x * theta\n",
        "\n",
        "theta = jax.random.normal(jax.random.PRNGKey(42))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wUldd3_2hObs"
      },
      "source": [
        "For the internal optimizer, which will fit $\\theta$, we will use RMSProp. For the external optimizer, which we will use to fit the learning rate, we will use Adam. In optax, we must use `optax.inject_hyperparams` in order to allow the outer optimizer to modify the learning rate of the inner optimizer.\n",
        "\n",
        "The inner optimizer is initialized with a learning rate of 0.1, and the outer optimizer is initialized with a learning rate of 0.03."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cWP2HN_KqQ6o"
      },
      "outputs": [],
      "source": [
        "init_learning_rate = jnp.array(0.1)\n",
        "meta_learning_rate = jnp.array(0.03)\n",
        "\n",
        "opt = optax.inject_hyperparams(optax.rmsprop)(learning_rate=init_learning_rate)\n",
        "meta_opt = optax.adam(learning_rate=meta_learning_rate)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uFG9tczunv3H"
      },
      "source": [
        "In the following code, we implement a step of gradient descent using the computed loss."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "suvw4J9Rn9Ju"
      },
      "outputs": [],
      "source": [
        "def loss(theta, x, y):\n",
        "  return optax.l2_loss(y, f(theta, x))\n",
        "\n",
        "\n",
        "def step(theta, state, x, y):\n",
        "  grad = jax.grad(loss)(theta, x, y)\n",
        "  updates, state = opt.update(grad, state)\n",
        "  theta = optax.apply_updates(theta, updates)\n",
        "  return theta, state"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cTbJC8h1oGe6"
      },
      "source": [
        "For the meta-learning part of the problem, we will use the inner step to compute an _outer_ loss value, and an _outer_ step.\n",
        "\n",
        "These look very similar to the inner step, however, this outer step will inject the value of the learning rate hyperparameter into the state for the inner optimizer.\n",
        "\n",
        "We use `jax.jit` on the outer step of our computation, to use JAX to optimize our computation."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2KojrF7FiJYD"
      },
      "outputs": [],
      "source": [
        "def outer_loss(eta, theta, state, samples):\n",
        "  state.hyperparams['learning_rate'] = jax.nn.sigmoid(eta)\n",
        "\n",
        "  for x, y in samples[:-1]:\n",
        "    theta, state = step(theta, state, x, y)\n",
        "\n",
        "  x, y = samples[-1]\n",
        "\n",
        "  return loss(theta, x, y), (theta, state)\n",
        "\n",
        "\n",
        "@jax.jit\n",
        "def outer_step(eta, theta, meta_state, state, samples):\n",
        "  grad, (theta, state) = jax.grad(\n",
        "      outer_loss, has_aux=True)(eta, theta, state, samples)\n",
        "\n",
        "  meta_updates, meta_state = meta_opt.update(grad, meta_state)\n",
        "  eta = optax.apply_updates(eta, meta_updates)\n",
        "\n",
        "  return eta, theta, meta_state, state"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wr7oTKoxoVSE"
      },
      "source": [
        "In the following, we put all of the code above together in order to fit a value for $\\theta$."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "N69sYlo_oUTZ"
      },
      "outputs": [],
      "source": [
        "state = opt.init(theta)\n",
        "# inverse sigmoid, to match the value we initialized the inner optimizer with.\n",
        "eta = -np.log(1. / init_learning_rate - 1)\n",
        "meta_state = meta_opt.init(eta)\n",
        "\n",
        "N = 7\n",
        "learning_rates = []\n",
        "thetas = []\n",
        "\n",
        "for i in range(2000):\n",
        "  samples = [next(g) for i in range(N)]\n",
        "  eta, theta, meta_state, state = outer_step(eta, theta, meta_state, state, samples)\n",
        "  learning_rates.append(jax.nn.sigmoid(eta))\n",
        "  thetas.append(theta)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lp0e2sS8pIqL"
      },
      "source": [
        "We can now plot the learning rates and values for $\\theta$ that we computed during our optimization,"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Gp45HrNyJuW3"
      },
      "outputs": [],
      "source": [
        "fig, (ax1, ax2) = plt.subplots(2);\n",
        "fig.suptitle('Meta-learning RMSProp\\'s learning rate');\n",
        "plt.xlabel('Step');\n",
        "\n",
        "ax1.semilogy(range(len(learning_rates)), learning_rates);\n",
        "ax1.set(ylabel='Learning rate');\n",
        "ax1.label_outer();\n",
        "\n",
        "plt.xlabel('Number of updates');\n",
        "ax2.semilogy(range(len(thetas)), thetas);\n",
        "\n",
        "ax2.label_outer();\n",
        "ax2.set(ylabel='Theta');"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tNvaCoeYpQRL"
      },
      "source": [
        "The profile of the learning rate over time seems fairly reasonable, it decays towards zero to help convergence."
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "name": "Meta-Learning",
      "private_outputs": true,
      "provenance": [
        {
          "file_id": "1UF4UGTpZvC0AWbko-4IFzsG43RH4b6WC",
          "timestamp": 1637243991785
        },
        {
          "file_id": "1ZRN4T0N8o-OvX4jIN52Beo3kzpNDM-Mx",
          "timestamp": 1637168142510
        }
      ]
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
